import tensorflow as tf
import tensorflow.keras as keras
from python_ai.common.xcommon import *
from tensorflow.keras import layers, activations

sep('load')
LEN_DICT = 1000
(x_train, y_train), (x_test, y_test) = keras.datasets.imdb.load_data(num_words=LEN_DICT)
check_shape_stdout(x_train, 'x_train')
check_shape_stdout(y_train, 'y_train')
check_shape_stdout(x_test, 'x_test')
check_shape_stdout(y_test, 'y_test')

sep('select')
N_SELECT = 1000
x_train = x_train[:N_SELECT]
y_train = y_train[:N_SELECT]

sep('pad sequence')
N_STEPS = 80
x_train = tf.keras.preprocessing.sequence.pad_sequences(x_train, N_STEPS)
check_shape_stdout(x_train, 'x_train')

sep('embed')
N_EMBEDDING = 300
x_train = layers.Embedding(LEN_DICT, N_EMBEDDING)(x_train, training=True)
check_shape_stdout(x_train, 'x_train')

sep('lstm 1')
N_RNN_HIDDEN = 128
x_train = layers.LSTM(N_RNN_HIDDEN,
                      return_sequences=True,
                      unroll=True,
                      dropout=0.2)(x_train, training=True)
check_shape_stdout(x_train, 'x_train')

sep('lstm 2')
x_train = layers.LSTM(N_RNN_HIDDEN,
                      return_sequences=False,
                      unroll=True,
                      dropout=0.2)(x_train, training=True)
check_shape_stdout(x_train, 'x_train')

sep('fc')
x_train = layers.Dense(1, activation=activations.sigmoid)(x_train)
check_shape_stdout(x_train, 'x_train')
